package com.xcj.springai.config;


import com.alibaba.cloud.ai.graph.*;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import com.alibaba.cloud.ai.graph.action.AsyncEdgeAction;
import com.xcj.springai.node.ExpanderNode;
import com.xcj.springai.node.HumanFeedbackNode;
import com.xcj.springai.node.SummaryNode;
import com.xcj.springai.node.TranslateNode;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.Map;

import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;


/**
 * 配置类
 * 作用： 配置并构建整个状态图（StateGraph）。
 * 职责：
 * 定义状态变量（如 query, expander_content, feedback 等）和替换策略；
 * 注册多个节点（expander, humanfeedback, translate）；
 * 配置节点间的跳转逻辑，包括普通边和条件边；
 * 打印 PlantUML 图用于可视化调试。
 * 角色： 整个状态图的“主装配器”
 */

@Configuration
public class GraphHumanConfiguration {
    private static final Logger logger = LoggerFactory.getLogger(GraphHumanConfiguration.class);

    @Bean
    public StateGraph humanGraph(ChatClient.Builder chatClientBuilder) throws GraphStateException {


        // 创建状态工厂
        OverAllStateFactory stateFactory = () -> {
            // 创建 OverAllState 对象
            OverAllState state = new OverAllState();
            // 注册键值对和替换策略
            state.registerKeyAndStrategy("query", new ReplaceStrategy());
            state.registerKeyAndStrategy("expander_number", new ReplaceStrategy());
            state.registerKeyAndStrategy("expander_content", new ReplaceStrategy());
            state.registerKeyAndStrategy("feedback", new ReplaceStrategy());
            state.registerKeyAndStrategy("threadId", new ReplaceStrategy());
            state.registerKeyAndStrategy("summary", new ReplaceStrategy());
            state.registerKeyAndStrategy("humanfeedback", new ReplaceStrategy());
            state.registerKeyAndStrategy("feedback", new ReplaceStrategy());
            state.registerKeyAndStrategy("queryVariants", new ReplaceStrategy());

            // 注册更多键值对和替换策略
            state.registerKeyAndStrategy("humannextnode", new ReplaceStrategy());
            state.registerKeyAndStrategy("translatelanguage", new ReplaceStrategy());
            state.registerKeyAndStrategy("translatecontent", new ReplaceStrategy());
            state.registerKeyAndStrategy("summary_content", new ReplaceStrategy());
            return state;
        };

        // 创建 StateGraph 对象并添加节点和边
        StateGraph stateGraph = new StateGraph(stateFactory)
                // 添加异步节点
                //扩展变体
                .addNode("expander", node_async(new ExpanderNode(chatClientBuilder)))
                //翻译变体
                .addNode("translate", node_async(new TranslateNode(chatClientBuilder)))
                //反馈变体
                .addNode("humanfeedback", node_async(new HumanFeedbackNode()))
                .addNode("summary", node_async(new SummaryNode()))

                // 添加边
                .addEdge(StateGraph.START, "expander")
                .addEdge("expander", "humanfeedback")
                // 添加条件边  这边添加条件分支的条件 可以有多个 然后是根据 HumanFeedbackDispatcher 的结果选择下一个节点
//                .addConditionalEdges("humanfeedback", new HumanFeedbackDispatcher(), Map.of(
                .addConditionalEdges("humanfeedback", AsyncEdgeAction.edge_async((new HumanFeedbackDispatcher())), Map.of(
                        "translate", "translate",
                        "summary", "summary",
//                        "test", "test",  //示例用的
                        StateGraph.END, StateGraph.END))

                // 添加另一条边  执行完 translate 后再执行

                .addEdge("translate",  StateGraph.END)
                .addEdge("summary", StateGraph.END);
        // 打印状态图
        logger.info("\n=== expander stateGraph ===");

        // 添加 PlantUML 打印
        // 获取状态图的 PlantUML 表示
        GraphRepresentation representation = stateGraph.getGraph(GraphRepresentation.Type.PLANTUML,
                "human flow");
        // 打印 UML 流

        logger.info("\n=== expander UML Flow ===");
        logger.info(representation.content());
        logger.info("==================================\n");

        return stateGraph;
    }

}
